#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Geppetto WebSocket 服务器
用于前后端通信
"""

import json
import asyncio
import logging
from typing import Dict, Any, Set
from tornado import web, websocket, ioloop
from netpyne_model import SimpleNetworkModel

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class GeppettoWebSocketHandler(websocket.WebSocketHandler):
    """Geppetto WebSocket 处理器"""
    
    # 存储所有连接的客户端
    clients: Set['GeppettoWebSocketHandler'] = set()
    
    # NetPyNE 模型实例
    model: SimpleNetworkModel = None
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client_id = None
    
    def check_origin(self, origin):
        """允许跨域连接"""
        return True
    
    def open(self):
        """WebSocket 连接打开"""
        self.client_id = id(self)
        GeppettoWebSocketHandler.clients.add(self)
        logger.info(f"客户端连接: {self.client_id}, 当前连接数: {len(self.clients)}")
        
        # 发送欢迎消息
        self.send_message({
            'type': 'connection',
            'status': 'connected',
            'message': 'Geppetto WebSocket 服务器已连接',
            'clientId': self.client_id
        })
    
    def on_message(self, message):
        """接收消息"""
        try:
            data = json.loads(message)
            logger.info(f"收到消息: {data.get('type', 'unknown')}")
            
            # 处理不同类型的消息
            msg_type = data.get('type')
            
            if msg_type == 'ping':
                self.handle_ping(data)
            elif msg_type == 'init_model':
                self.handle_init_model(data)
            elif msg_type == 'get_network':
                self.handle_get_network(data)
            elif msg_type == 'run_simulation':
                self.handle_run_simulation(data)
            elif msg_type == 'get_simulation_data':
                self.handle_get_simulation_data(data)
            elif msg_type == 'execute_code':
                self.handle_execute_code(data)
            else:
                self.send_error(f"未知的消息类型: {msg_type}")
        
        except json.JSONDecodeError as e:
            logger.error(f"JSON 解析错误: {e}")
            self.send_error(f"JSON 解析错误: {str(e)}")
        except Exception as e:
            logger.error(f"处理消息错误: {e}", exc_info=True)
            self.send_error(f"处理消息错误: {str(e)}")
    
    def on_close(self):
        """WebSocket 连接关闭"""
        GeppettoWebSocketHandler.clients.discard(self)
        logger.info(f"客户端断开: {self.client_id}, 当前连接数: {len(self.clients)}")
    
    def send_message(self, data: Dict[str, Any]):
        """发送消息给客户端"""
        try:
            message = json.dumps(data)
            self.write_message(message)
        except Exception as e:
            logger.error(f"发送消息错误: {e}")
    
    def send_error(self, error_message: str):
        """发送错误消息"""
        self.send_message({
            'type': 'error',
            'error': error_message
        })
    
    def handle_ping(self, data: Dict[str, Any]):
        """处理 ping 消息"""
        self.send_message({
            'type': 'pong',
            'timestamp': data.get('timestamp')
        })
    
    def handle_init_model(self, data: Dict[str, Any]):
        """初始化 NetPyNE 模型"""
        try:
            logger.info("初始化 NetPyNE 模型...")
            
            # 创建模型
            if GeppettoWebSocketHandler.model is None:
                GeppettoWebSocketHandler.model = SimpleNetworkModel()
            
            self.send_message({
                'type': 'model_initialized',
                'status': 'success',
                'message': 'NetPyNE 模型初始化成功'
            })
            
            logger.info("NetPyNE 模型初始化成功")
        
        except Exception as e:
            logger.error(f"初始化模型错误: {e}", exc_info=True)
            self.send_error(f"初始化模型错误: {str(e)}")
    
    def handle_get_network(self, data: Dict[str, Any]):
        """获取网络数据"""
        try:
            logger.info("获取网络数据...")
            
            # 确保模型已初始化
            if GeppettoWebSocketHandler.model is None:
                self.handle_init_model(data)
            
            # 创建网络
            GeppettoWebSocketHandler.model.create_network()
            
            # 获取网络数据
            network_data = GeppettoWebSocketHandler.model.get_network_data()
            
            self.send_message({
                'type': 'network_data',
                'status': 'success',
                'data': network_data
            })
            
            logger.info(f"网络数据发送成功: {len(network_data['cells'])} 个细胞")
        
        except Exception as e:
            logger.error(f"获取网络数据错误: {e}", exc_info=True)
            self.send_error(f"获取网络数据错误: {str(e)}")
    
    def handle_run_simulation(self, data: Dict[str, Any]):
        """运行仿真"""
        try:
            logger.info("运行仿真...")
            
            # 确保模型已初始化
            if GeppettoWebSocketHandler.model is None:
                self.handle_init_model(data)
            
            # 发送仿真开始消息
            self.send_message({
                'type': 'simulation_started',
                'status': 'running',
                'message': '仿真开始运行...'
            })
            
            # 运行仿真
            sim_data = GeppettoWebSocketHandler.model.run_simulation()
            
            # 发送仿真完成消息
            self.send_message({
                'type': 'simulation_completed',
                'status': 'success',
                'message': '仿真运行完成'
            })
            
            logger.info("仿真运行完成")
        
        except Exception as e:
            logger.error(f"运行仿真错误: {e}", exc_info=True)
            self.send_error(f"运行仿真错误: {str(e)}")
    
    def handle_get_simulation_data(self, data: Dict[str, Any]):
        """获取仿真数据"""
        try:
            logger.info("获取仿真数据...")
            
            # 确保模型已初始化
            if GeppettoWebSocketHandler.model is None:
                self.send_error("模型未初始化，请先运行仿真")
                return
            
            # 获取仿真数据
            simulation_data = GeppettoWebSocketHandler.model.get_simulation_data()
            
            if simulation_data is None:
                self.send_error("仿真数据不可用，请先运行仿真")
                return
            
            self.send_message({
                'type': 'simulation_data',
                'status': 'success',
                'data': simulation_data
            })
            
            logger.info(f"仿真数据发送成功: {len(simulation_data['spikes'])} 个 spikes")
        
        except Exception as e:
            logger.error(f"获取仿真数据错误: {e}", exc_info=True)
            self.send_error(f"获取仿真数据错误: {str(e)}")
    
    def handle_execute_code(self, data: Dict[str, Any]):
        """执行 Python 代码"""
        try:
            code = data.get('code', '')
            if not code:
                self.send_error("代码为空")
                return
            
            logger.info(f"执行代码: {code[:100]}...")
            
            # 创建执行环境
            exec_globals = {
                'model': GeppettoWebSocketHandler.model,
                'SimpleNetworkModel': SimpleNetworkModel
            }
            exec_locals = {}
            
            # 执行代码
            exec(code, exec_globals, exec_locals)
            
            # 获取结果
            result = exec_locals.get('result', 'Code executed successfully')
            
            self.send_message({
                'type': 'code_executed',
                'status': 'success',
                'result': str(result)
            })
            
            logger.info("代码执行成功")
        
        except Exception as e:
            logger.error(f"执行代码错误: {e}", exc_info=True)
            self.send_error(f"执行代码错误: {str(e)}")
    
    @classmethod
    def broadcast(cls, message: Dict[str, Any]):
        """广播消息给所有客户端"""
        for client in cls.clients:
            try:
                client.send_message(message)
            except Exception as e:
                logger.error(f"广播消息错误: {e}")


class HealthCheckHandler(web.RequestHandler):
    """健康检查处理器"""
    
    def get(self):
        self.set_header('Content-Type', 'application/json')
        self.write({
            'status': 'ok',
            'service': 'Geppetto WebSocket Server',
            'clients': len(GeppettoWebSocketHandler.clients)
        })


def make_app():
    """创建 Tornado 应用"""
    return web.Application([
        (r'/geppetto', GeppettoWebSocketHandler),
        (r'/health', HealthCheckHandler),
    ])


def start_server(port=8888):
    """启动服务器"""
    app = make_app()
    app.listen(port)
    logger.info(f"Geppetto WebSocket 服务器启动在端口 {port}")
    logger.info(f"WebSocket URL: ws://localhost:{port}/geppetto")
    logger.info(f"Health Check URL: http://localhost:{port}/health")
    
    try:
        ioloop.IOLoop.current().start()
    except KeyboardInterrupt:
        logger.info("服务器关闭")


if __name__ == '__main__':
    start_server(port=8888)

